"""
PlumeSoft 量化代码库
作者：PlumeSoft
用途：贝叶斯选币炼丹炉
"""

from datetime import datetime
import os
import shutil
import pandas as pd
from dateutil.relativedelta import relativedelta
from program import Config as Cfg
from program.plumesoft_bys.bys_back_test import batch_run_stg_back_test
from program.plumesoft_bys.bys_batch import (
    Batch_Bayesian_Search,
    get_back_test_eval_score_dict,
    run_top_stg_back_test,
)
import program.plumesoft_bys.bys_function as func
from program.plumesoft_bys.bys_common import (
    add_code_to_stg_file,
    init_env,
    log,
    dbg,
)


def Rolling_Bayesian_Search(params: dict, bys):
    init_env(params)

    all_stime = datetime.now()
    base_stg_folder = params["base_stg_folder"]
    base_stg_file = params["base_stg_file"]

    log(
        f"开始滚动贝叶斯炼丹，基础策略文件：{base_stg_folder}.{base_stg_file}",
        send_wx=True,
    )
    dbg(f"参数: {params}")

    a_df = pd.DataFrame()
    work_path = os.path.join(
        params["data_path"],
        func.replace_macro_path(params, params["rolling_work_path"]),
    )
    study_rolling_csv = os.path.join(
        work_path, func.replace_macro_path(params, params["study_rolling_csv"])
    )
    if os.path.isfile(study_rolling_csv):
        a_df = pd.read_csv(
            study_rolling_csv,
            encoding="gbk",
            parse_dates=[
                "样本内起始",
                "样本内结束",
                "样本外起始",
                "样本外结束",
                "实盘起始",
                "实盘结束",
            ],
        )
    if a_df.empty:
        rolling_id = datetime.now().strftime("%Y%m%d$%H%M%S")
    else:
        rolling_id = a_df["滚动号"].iloc[0]
    params["滚动号"] = rolling_id

    # 输出的滚动策略文件
    rolling_dst_fn = os.path.join(
        Cfg.root_path,
        "program",
        params["output_stg_folder"],
        f"{base_stg_file}$滚动${rolling_id}.py",
    )
    os.makedirs(os.path.dirname(rolling_dst_fn), exist_ok=True)

    # 取得用于计算评分的参数
    cls = __import__(f"program.{base_stg_folder}.{base_stg_file}", fromlist=("",))
    stg_params = {}
    for key in params:
        if hasattr(cls, key) and getattr(cls, key) is not None:
            params[key] = getattr(cls, key)
            stg_params[key] = params[key]
    log(f"从策略文件中更新配置参数: {str(stg_params)}")

    # 计算滚动结束日期
    rolling_date_end = pd.to_datetime(params["rolling_date_end"])
    if rolling_date_end is None:
        rolling_date_end = pd.to_datetime(datetime.today().date())

    rolling_date_start = pd.to_datetime(params["rolling_date_start"])
    rolling_sample_months_min = params["rolling_sample_months_min"]
    rolling_sample_months_max = params["rolling_sample_months_max"]
    if (
        rolling_sample_months_max < rolling_sample_months_min
        or rolling_sample_months_min < 1
    ):
        log(
            f"rolling_sample_months_max 必须大于 rolling_sample_months_min，且两者都必须大于 0"
        )
        exit()

    rolling_back_test_months = params["rolling_back_test_months"]
    if rolling_back_test_months < 0:
        log(f"rolling_back_test_months 必须大于或等于 0")
        exit()

    rolling_trade_months = params["rolling_trade_months"]
    if rolling_trade_months < 1:
        log(f"rolling_trade_months 必须大于 0")
        exit()

    if (
        rolling_date_start
        + relativedelta(months=rolling_sample_months_min + rolling_back_test_months + rolling_trade_months)
        >= rolling_date_end
    ):
        log(f"滚动开始日期加样本内、样本外和模拟实盘月份后必须小于滚动结束日期，请检查滚动日期参数")
        exit()

    rolling_back_test_in_front = params["rolling_back_test_in_front"]

    # 计算最开始的实盘起始日期
    curr_trade_date_start = pd.to_datetime(
        rolling_date_start
        + relativedelta(months=rolling_sample_months_min + rolling_back_test_months)
    )
    round_id = 1
    while curr_trade_date_start < rolling_date_end:
        # 实盘结束日期
        curr_trade_date_end = min(
            pd.to_datetime(rolling_date_end + relativedelta(days=1)),
            pd.to_datetime(
                curr_trade_date_start + relativedelta(months=rolling_trade_months)
            ),
        )
        # 样本外回测结束日期，等于实盘起始日期
        curr_back_test_date_end = curr_trade_date_start
        # 样本外回测起始日期
        curr_back_test_date_start = pd.to_datetime(
            curr_back_test_date_end - relativedelta(months=rolling_back_test_months)
        )
        # 样本内回测结束日期，等于样本外起始日期
        curr_sample_date_end = curr_back_test_date_start
        if rolling_sample_months_max is None:
            curr_sample_date_start = rolling_date_start
        else:
            curr_sample_date_start = max(
                rolling_date_start,
                pd.to_datetime(
                    curr_sample_date_end
                    - relativedelta(months=rolling_sample_months_max)
                ),
            )

        # 样本外放在前面
        if rolling_back_test_in_front:
            curr_sample_date_end = curr_back_test_date_end
            curr_back_test_date_start = curr_sample_date_start
            curr_back_test_date_end = pd.to_datetime(
                curr_back_test_date_start
                + relativedelta(months=rolling_back_test_months)
            )
            curr_sample_date_start = curr_back_test_date_end

        # 结束日期向前推一天
        curr_sample_date_end = pd.to_datetime(
            curr_sample_date_end - relativedelta(days=1)
        )
        if rolling_back_test_months > 0:
            curr_back_test_date_end = pd.to_datetime(
                curr_back_test_date_end - relativedelta(days=1)
            )
        curr_trade_date_end = pd.to_datetime(
            curr_trade_date_end - relativedelta(days=1)
        )

        log("===========================================")
        msg = f"开始第 {round_id} 轮滚动炼丹...\n"
        msg += f"样本内区间：{curr_sample_date_start.date()} - {curr_sample_date_end.date()}\n"
        if rolling_back_test_months > 0:
            msg += f"样本外区间：{curr_back_test_date_start.date()} - {curr_back_test_date_end.date()}\n"
        else:
            curr_back_test_date_start = None
            curr_back_test_date_end = None
            msg += "样本外区间：禁用\n"
        msg += (
            f"实盘区间：{curr_trade_date_start.date()} - {curr_trade_date_end.date()}"
        )
        log(msg, send_wx=True)

        # 检查是否断点续传
        if not a_df.empty:
            ex_df = a_df[a_df["滚动轮次"] == round_id]
            if len(ex_df) > 0:
                if (
                    ex_df["实盘起始"].iloc[0] != curr_trade_date_start
                    or ex_df["实盘结束"].iloc[0] != curr_trade_date_end
                ):
                    log(
                        f"找到第 {round_id} 轮滚动炼丹数据与计算结果不匹配，无法进行断点续传，请检查参数或重新开始滚动炼丹"
                    )
                    return
                log(f"找到第 {round_id} 轮已完成的滚动炼丹数据，跳过...")
                round_id += 1
                curr_trade_date_start = pd.to_datetime(
                    curr_trade_date_start + relativedelta(months=rolling_trade_months)
                )
                continue

        # 配置用于批量炼丹的参数
        new_params = {}
        new_params["sample_date_start"] = curr_sample_date_start
        new_params["sample_date_end"] = curr_sample_date_end
        new_params["use_back_test_mode"] = rolling_back_test_months > 0
        new_params["back_test_date_start"] = curr_back_test_date_start
        new_params["back_test_date_end"] = curr_back_test_date_end
        new_params["output_stg_batch_back_test_mode"] = None
        new_params["dashboard_enabled"] = False
        new_params["batch_work_path"] = params["rolling_work_path"]
        new_params["batch_finished_path"] = os.path.join(
            params["rolling_finished_path"],
            f"批量_{curr_trade_date_start.strftime('%Y%m%d')}-{curr_trade_date_end.strftime('%Y%m%d')}/轮次_$ROUND-ID$",
        )
        # 滚动炼丹中配置的参数不允许被策略中的定义修改
        ignore_params = list(new_params.keys())
        curr_params = params.copy()
        curr_params.update(new_params)

        # 执行一轮批量炼丹
        df = Batch_Bayesian_Search(curr_params, bys, ignore_params=ignore_params)

        # 保存到滚动结果中的数据
        res_dt = {
            "滚动号": rolling_id,
            "滚动轮次": round_id,
            "批次": curr_params["批次"],
            "样本内起始": curr_sample_date_start,
            "样本内结束": curr_sample_date_end,
            "样本外起始": curr_back_test_date_start,
            "样本外结束": curr_back_test_date_end,
            "实盘起始": curr_trade_date_start,
            "实盘结束": curr_trade_date_end,
            "最高得分": df["score"].iloc[0],
            "策略名称": f"子策略_{round_id:03}",
        }
        a_df = pd.concat([a_df, pd.DataFrame([res_dt])], ignore_index=True)
        a_df.to_csv(study_rolling_csv, index=False, encoding="gbk")

        # 生成所有子策略的信息
        stg_dict = {}
        for rec in a_df.to_records():
            stg_dict[rec["策略名称"]] = {
                "起始日期": str(pd.to_datetime(rec["实盘起始"]).date()),
                "结束日期": str(pd.to_datetime(rec["实盘结束"]).date()),
            }

        code = f"\n# 炼丹时间: {datetime.now()}\n"
        code += f'# 基础策略: {params["base_stg_file"]}\n'
        str_stg_dict = func.dict_to_json_str(stg_dict)
        code += f"bys_rolling_stg_dict = {str_stg_dict}\n\n"
        add_code_to_stg_file(
            os.path.join(os.path.dirname(__file__), "bys_rolling_template.py"),
            rolling_dst_fn,
            "# @$开始定义滚动炼丹结果数据$@#",
            "# @$结束定义滚动炼丹结果数据$@#",
            code,
            append_mode=False,
        )
        log(f'已生成滚动策略文件，保存在 {params["output_stg_folder"]} 文件夹中。')

        # 取得批量炼丹排名第一的策略文件名
        res_stg_folder = os.path.join(
            Cfg.root_path, "program", params["output_stg_folder"]
        )
        res_stg_file = os.path.join(
            res_stg_folder,
            f"{base_stg_file}$批量${df['批次'].iloc[0]}$001.py",
        )
        # 复制到滚动策略子目录下
        dst_fn = os.path.join(
            rolling_dst_fn.replace(".py", "_子策略"), res_dt["策略名称"] + ".py"
        )
        os.makedirs(os.path.dirname(dst_fn), exist_ok=True)
        shutil.move(res_stg_file, dst_fn)

        # 删除不需要的批量策略文件
        for file in os.listdir(res_stg_folder):
            if file.endswith(".py") and file.startswith(
                f"{base_stg_file}$批量${df['批次'].iloc[0]}$"
            ):
                os.remove(os.path.join(res_stg_folder, file))

        # 根据 output_stg_top_num_rolling 的定义生成模拟实盘的回测数据
        study_rolling_back_test_csv = os.path.join(
            work_path,
            func.replace_macro_path(params, params["study_rolling_back_test_csv"]),
        )
        if (
            params["output_stg_top_num_rolling"] is not None
            and params["output_stg_top_num_rolling"] > 0
        ):
            # 生成临时策略文件用于计算样本外表现，生成最终的评分
            df = df.head(params["output_stg_top_num_rolling"]).reset_index(drop=True)
            df["排名"] = df.index + 1
            df["滚动轮次"] = round_id
            df["样本内起始"] = curr_sample_date_start
            df["样本内结束"] = curr_sample_date_end
            df["样本外起始"] = curr_back_test_date_start
            df["样本外结束"] = curr_back_test_date_end
            df["模拟实盘起始"] = curr_trade_date_start
            df["模拟实盘结束"] = curr_trade_date_end

            # 生成临时策略文件用于计算样本外表现，生成最终的评分
            df = run_top_stg_back_test(
                params,
                bys,
                df,
                cls,
                "模拟实盘",
                f"模拟实盘 {curr_trade_date_start} - {curr_trade_date_end}",
                curr_trade_date_start,
                curr_trade_date_end,
                get_back_test_eval_score_dict(params),
                update_score=False,
            )

            sorted_cols = [
                "滚动轮次",
                "排名",
                "批次",
                "轮次",
                "迭代",
                "过滤参数",
                "用户参数",
                "排序参数",
                "高级因子",
                "样本内起始",
                "样本内结束",
                "回测数据",
                "炼丹评分",
                "样本外起始",
                "样本外结束",
                "样本外回测数据",
                "样本外评分",
                "模拟实盘起始",
                "模拟实盘结束",
                "模拟实盘回测数据",
                "模拟实盘评分",
            ]
            sorted_cols = [col for col in sorted_cols if col in df.columns]
            df = df[sorted_cols]
            if os.path.isfile(study_rolling_back_test_csv):
                df.to_csv(
                    study_rolling_back_test_csv,
                    index=False,
                    mode="a",
                    header=False,
                    encoding="gbk",
                )
            else:
                df.to_csv(
                    study_rolling_back_test_csv, index=False, mode="w", encoding="gbk"
                )

        round_id += 1
        curr_trade_date_start = pd.to_datetime(
            curr_trade_date_start + relativedelta(months=rolling_trade_months)
        )

    # 全部完成自动回测
    if params["output_stg_rolling_auto_back_test"]:
        log(
            f"开始使用框架 Config.py 中配置的参数执行策略回测，回测区间：{Cfg.start_date} - {Cfg.end_date}"
        )
        batch_run_stg_back_test(
            params["output_stg_folder"],
            [os.path.basename(rolling_dst_fn).replace(".py", "")],
            Cfg.start_date,
            Cfg.end_date,
            if_plot=True,
            pic_dir=params["pic_path"],
            open_pic=True,
        )

    # 取得滚动目录
    finished_path = os.path.dirname(
        os.path.join(
            params["data_path"],
            func.replace_macro_path(params, params["rolling_finished_path"]),
        )
    )

    # 复制结果策略到结果文件夹
    out_stg_path = os.path.join(finished_path, "结果策略")
    os.makedirs(out_stg_path, exist_ok=True)
    shutil.copy(
        rolling_dst_fn, os.path.join(out_stg_path, os.path.basename(rolling_dst_fn))
    )
    sub_stg_path = rolling_dst_fn.replace(".py", "_子策略")
    shutil.copytree(
        sub_stg_path,
        os.path.join(out_stg_path, os.path.basename(sub_stg_path)),
        dirs_exist_ok=True,
    )

    for file in [study_rolling_csv, study_rolling_back_test_csv]:
        # 移动滚动结果文件到结果文件夹，结束滚动炼丹
        if os.path.isfile(file):
            shutil.move(
                file,
                os.path.join(finished_path, os.path.basename(file)),
            )

    log(f"滚动炼丹完成，总用时 {datetime.now() - all_stime}", send_wx=True)

    return a_df
